# Specify parameters for prior
#mutable struct priors
#    mu_phi::Array{Float64,2}
#    P_phi::Array{Float64,2}
#    nu_iwish::Int64
#    sig_iwish::Array{Float64,2}
#end

# Specify parameters for prior
mutable struct priors_ACP
    beta0::Array{Float64,2}
    alp0::Array{Float64,2}
    Vbeta::Array{Float64,2}
    Valp::Array{Float64,2}
    nu::Array{Float64,2}
    S::Array{Float64,2}
end

##

function data_to_YY(data, T0, nlags)

    # Data handling
    Tn    = size(data)[1]
    YY    = data[T0+1:Tn,:]
    T     = size(YY)[1] # sample size

    XX    = data[T0:Tn-1,:]
    if nlags > 1
        for  ii = 2:nlags
            XX = [ XX data[T0+1-ii:Tn-ii,:]]
        end
    end

    # OLS estimator
    PHIhat   = \((XX'XX),(XX'YY))
    Shat     = (YY'*YY) - (YY'*XX)*PHIhat
    SIGMAhat = Shat/T

return YY, XX, PHIhat, Shat, SIGMAhat
end

##
function compdata_to_WW(agg_data, cross_data, nlags)

    # Data handling
    Tn    = size(agg_data)[1]

    # First data row is to initialize lag for obs nlags+1
    # Thus, we start at nlags and we go to end of sample Tn

    agg_XX    = agg_data[nlags:Tn,:]
    if nlags > 1
        for  ii = 2:nlags
            agg_XX = [ agg_XX agg_data[nlags+1-ii:Tn+1-ii,:]]
        end
    end

    cross_XX    = cross_data[nlags:Tn,:]
    if nlags > 1
        for  ii = 2:nlags
            cross_XX = [ cross_XX cross_data[nlags+1-ii:Tn+1-ii,:]]
        end
    end

   data_YYcross = [agg_XX cross_XX]

return data_YYcross
end



function prior_ACPi(n,p,var_i,lambda,sig2,n_agg,igammadf)
    ki = var_i + n*p - 1
    mi = zeros(ki,1)
    Vi = zeros(ki,1)

    # construct Vi
    for j=1:ki
        if j<= n*p
            l = ceil(j/n) # lag length
            idx = mod(j,n) # variable index for j
            if idx==0
                idx = n
            end
        else
                idx = j-n*p
        end

        # Old
        # if j > n*p # alpha_i
        #     Vi[j] = 1/sig2[idx]
        # elseif idx == var_i # own lag
        #     Vi[j]=1/lambda[1]*(1/(sig2[idx]*l^lambda[4]))
        # elseif (idx > n_agg && idx != var_i) # var i belongs to Y and j belongs to a
        #     Vi[j]=1/lambda[1]*(1/(lambda[2]*sig2[idx]*l^lambda[4]))
        # elseif (idx <= n_agg && idx != var_i) # var i belongs to a and j belongs to Y
        #     Vi[j]=1/lambda[1]*(1/(lambda[3]*sig2[idx]*l^lambda[4]))
        # end
        if j > n*p # alpha_i
            Vi[j] = 1/sig2[idx]
        elseif (var_i <= n_agg && idx > n_agg ) # var i belongs to Y and regressor j belongs to a
            Vi[j]=1/lambda[1]*(1/(lambda[2]*sig2[idx]*l^lambda[4]))
        elseif (var_i > n_agg && idx <= n_agg ) # var i belongs to a and regressor j belongs to Y
            Vi[j]=1/lambda[1]*(1/(lambda[3]*sig2[idx]*l^lambda[4]))
        else # same group
            Vi[j]=1/lambda[1]*(1/(sig2[idx]*l^lambda[4]))
        end
end

Si = sig2[var_i]/2
nui = (igammadf+var_i-n)/2#1+var_i/2

return mi, Vi, nui, Si

end

function prior_ACP_redu(YY, SIGMAhat, n_agg, igammadf, nlags, lambda)
    # elicits the asymmetric conjugate prior on the reduced-form parameterization

    T = size(YY)[1]
    n = size(YY)[2]
    sig2 = diag(SIGMAhat)

    p = nlags
    n_cross = n-n_agg

    k_beta = n*n*p
    k_alp = Int(n*(n-1)/2)

    beta0 = zeros(Int(k_beta/n),n)
    alp0  = zeros(k_alp,1)
    Vbeta = zeros(Int(k_beta/n),n)
    Valp  = zeros(k_alp,1)
    nu    = zeros(n,1)
    S     = zeros(n,1)

    count_alp = 0

    for ii=1:n

        mi, Vi, nui, Si = prior_ACPi(n,p,ii,lambda,sig2,n_agg,igammadf)
        beta0[:,ii]     = mi[1:Int(k_beta/n)]
        alp0[(count_alp+1):(count_alp+ii-1)] = mi[Int(k_beta/n)+1:end]
        Vbeta[:,ii]     = Vi[1:Int(k_beta/n)]
        Valp[(count_alp+1):(count_alp+ii-1)] = Vi[Int(k_beta/n)+1:end]
        nu[ii] = nui
        S[ii] = Si
        count_alp = count_alp+ii-1

    end

    prior = priors_ACP(beta0,alp0,Vbeta,Valp,nu,S)

return prior

end


function prior_ACP_stru(YY, SIGMAhat, n_agg, igammadf, nlags, lambda)
# elicits the asymmetric conjugate prior on the structural form parameterizaion

T = size(YY)[1]
n = size(YY)[2]
sig2 = diag(SIGMAhat)

p = nlags
n_cross = n-n_agg

k_beta = n*n*p
prior_stru = prior_ACP_redu(YY, SIGMAhat, n_agg, igammadf, nlags, lambda)
alp0  = prior_stru.alp0
beta0 = prior_stru.beta0
Valp  = prior_stru.Valp
nu    = prior_stru.nu
S     = prior_stru.S
Vbeta = zeros(Int(k_beta/n),n)

for ii = 1:n
    for jj=1:n*p
        if ii==1
            Vbeta[jj,ii] = prior_stru.Vbeta[jj,ii]
        else
            Vbeta[jj,ii] = prior_stru.Vbeta[jj,ii] + sum(prior_stru.Vbeta[jj,1:ii-1] + (prior_stru.beta0[jj,1:ii-1].^2)./sig2[1:ii-1])
        end
    end
end

prior = priors_ACP(beta0,alp0,Vbeta,Valp,nu,S)

return prior

end

function ml_VAR_ACP(nlags,YY,XX,prior)
# computes the marginal likelihood value under the asymmetric conjugate prior

    T = size(YY)[1]
    n = size(YY)[2]
    p = nlags

    lml = -n*T/2*log(2*pi)
    count_alp = 0

    for ii = 1:n
        yi = YY[:,ii]
        ki = n*p+ii-1
        mi = [prior.beta0[:,ii]; prior.alp0[count_alp+1:count_alp+ii-1]]
        Vi = sparse(1:ki, 1:ki, [prior.Vbeta[:,ii]; prior.Valp[count_alp+1:count_alp+ii-1]])
        nui = prior.nu[ii]
        Si = prior.S[ii]
        Xi = [XX -YY[:,1:ii-1]]

        iVi = Vi\sparse(I,ki,ki)
        Kthetai = iVi + Xi'*Xi
        CKthetai = cholesky(Kthetai).L
        thetai_hat = CKthetai'\(CKthetai\(iVi*mi + Xi'*yi))
        Si_hat = Si + (yi'*yi + mi'*iVi*mi - thetai_hat'*Kthetai*thetai_hat)/2

        lml = lml - 1/2*(sum(log.(diag(Vi))) + 2*sum(log.(diag(CKthetai)))) + nui*log(Si) - (nui+T/2)*log(Si_hat) + (logabsgamma(nui+T/2))[1] - (logabsgamma(nui))[1] # p29 of Appendix B in Joshua Chen (QE)

        count_alp = count_alp + ii -1

    end

    return lml

end

function sample_ThetaSig(YY,XX,PHIhat,Shat,SIGMAhat,nlags,prior,nsim)
# posterior draws of alpha, beta, Sigma under the asymmetric conjugate prior

    # YY, XX, PHIhat, Shat, SIGMAhat = data_to_YY(data, T0, nlags)
    T  = size(YY)[1]
    n  = size(YY)[2]
    p = nlags

    k_beta = n*n*p
    k_alp = n*(n-1)/2
    BBeta = zeros(nsim, k_beta)
    AAlp  = zeros(nsim, Int(k_alp))
    SSig  = zeros(nsim,n)
    count_alp = 0

    for ii=1:n
        yi = YY[:,ii]
        ki = n*p+ii-1
        mi = [prior.beta0[:,ii]; prior.alp0[count_alp+1:count_alp+ii-1]]
        Vi = sparse(1:ki, 1:ki, [prior.Vbeta[:,ii]; prior.Valp[count_alp+1:count_alp+ii-1]])
        nui = prior.nu[ii]
        Si = prior.S[ii]
        Xi = [XX -YY[:,1:ii-1]]

        # compute the paramters of the posterior distribution
        iVi = Vi\sparse(I,ki,ki)
        Kthetai = iVi + Xi'*Xi
        CKthetai = cholesky(Kthetai).L
        thetai_hat = CKthetai'\(CKthetai\(iVi*mi + Xi'*yi))
        Si_hat = Si + (yi'*yi + mi'*iVi*mi - thetai_hat'*Kthetai*thetai_hat)/2

        # sample sig and theta
        Sigi = 1 ./rand(Gamma(nui+T/2, 1/Si_hat),nsim)
        U = randn(nsim,ki).*repeat(sqrt.(Sigi),1,ki)
        Thetai = repeat(thetai_hat',nsim,1) + U/CKthetai

        SSig[:,ii] = Sigi
        BBeta[:,(ii-1)*n*p+1:ii*n*p] = Thetai[:,1:n*p]
        AAlp[:,count_alp+1:count_alp+ii-1] = Thetai[:,n*p+1:end]
        count_alp = count_alp+ii-1

    end

    return AAlp, BBeta, SSig

end

function getReducedForm(store_alp,store_beta,store_Sig)

    nsim = size(store_Sig)[1]
    n = size(store_Sig)[2]

    k_beta = size(store_beta)[2]
    p = Int((k_beta/n)/n)
    A_id = nonzeros(sparse(tril(reshape(1:n^2,n,n),-1)'))
    A = Matrix(1.0I,n,n)

    store_Btilde = zeros(nsim, n*n*p,1)
    store_Sigtilde = zeros(nsim,n,n)

    # compute reduced-form parameters
    for isim = 1:nsim
        aalp = store_alp[isim,:]
        bbeta = store_beta[isim,:]
        ssig = store_Sig[isim,:]

        # transform the parameters into reduced form
        A[A_id] = aalp
        Sig = (A\sparse(1:n,1:n,ssig))/A'
        Btilde = (A\(reshape(bbeta,n*p,n)'))'

        store_Btilde[isim,:] = Btilde[:]
        store_Sigtilde[isim,:,:] = Sig

    end

return store_Btilde, store_Sigtilde

end
##

function IRFredu(A,L,nstep,sh_size,sh_id)

    np = size(A)[1]
    n = size(A)[2]

    p = Int(np/n)

    response = zeros(n,nstep)
    Acomp = [A'; sparse(Array(1:n*(p-1)), Array(1:n*(p-1)), ones(n*(p-1)), n*(p-1), np)]
    response[:,1] = sh_size*L[:,sh_id]

    Apower = Acomp

    for hh = 2:nstep
        Apower = Apower*Acomp
#        response[:,hh] = Apower[1:n,1:n]*L[:,1]
        response[:,hh] = Apower[1:n,1:n]*response[:,1]
    end

return response

end

function IRFredu_maxGini(A,L,nstep,sh_size)

    np = size(A)[1]
    n = size(A)[2]

    p = Int(np/n)

    lenq = 500 # out of 500 rotations, find the one that maximizes the Gini coefficients
    #Random.seed!(100000*seedindx+10*t+3+seedoffset)
    qgrid = randn(lenq,n_cross)

    for qq = 1:lenq
       qgrid[qq,:] = qgrid[qq,:]./norm(qgrid[qq,:])
    end
    qgrid = [zeros(lenq,n_agg) qgrid]; # save candidate shocks

    response = zeros(n,nstep)
    Acomp = [A'; sparse(Array(1:n*(p-1)), Array(1:n*(p-1)), ones(n*(p-1)), n*(p-1), np)]

    # find the optimal q
    Gini_vec = zeros(lenq,1)

    for qq = 1:lenq
        temp_IRF = sh_size*L*qgrid[qq,:] # on impact
        PhatDensCoef_hh    = coefRecover(temp_IRF[n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
        PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,temp_IRF[3]+mean_unrate, minimum(xgrid), maximum(xgrid))
        PhatDens_IRF = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
        Gini_vec[qq] = GiniCoef(PhatDens_IRF, xgrid)
    end

    maxGini_qstar = qgrid[getindex.(findall(Gini_vec.-maximum(Gini_vec).==0),[1 2])[1],:]

    response[:,1] = -sh_size*L*maxGini_qstar

    Apower = Acomp

    for hh = 2:nstep
        Apower = Apower*Acomp
#        response[:,hh] = Apower[1:n,1:n]*L[:,1]
        response[:,hh] = Apower[1:n,1:n]*response[:,1]
    end

return response

end

function IRFredu_Anatomy_Agg(A,L,nstep,sh_size, select_agg)

    np = size(A)[1]
    n = size(A)[2]

    p = Int(np/n)

    lenq = 500 # out of 500 qs, find the one that maximizes the closed form variance
    #Random.seed!(100000*seedindx+10*t+3+seedoffset)
    qgrid = randn(lenq,n)

    for qq = 1:lenq
       qgrid[qq,:] = qgrid[qq,:]./norm(qgrid[qq,:])
    end

    response = zeros(n,nstep)
    Acomp = [A'; sparse(Array(1:n*(p-1)), Array(1:n*(p-1)), ones(n*(p-1)), n*(p-1), np)]

    M = zeros(np, n)
    M[1:n,1:n] = Matrix(1.0I,n,n)

    Var_Agg = zeros(lenq,1)
        for qq = 1:lenq
            temp_q = qgrid[qq,:]
            SIGMAtrq = Symmetric(M*L*temp_q*temp_q'*L'*M')
            SIGMAq = lyapd(Acomp, SIGMAtrq) # solve Lyaponov equation
            Var_Agg[qq] = diag(SIGMAq)[select_agg] # variance of aggregate variables
        end

    optimal_q = qgrid[findmax(vec(Var_Agg))[2],:]

     if (select_agg == 2 && (L*optimal_q)[select_agg] < 0) # on impact, targeted agg variable's response is positive (for GDP: select_agg = 2)
         optimal_q = -optimal_q
     end

    if (select_agg == 3 && (L*optimal_q)[select_agg] > 0)  # on impact, targeted agg variable's response is negative (for unemployment: select_agg = 3)
        optimal_q = -optimal_q
    end

    response[:,1] = sh_size*L*optimal_q

    Apower = Acomp

    for hh = 2:nstep
        Apower = Apower*Acomp
#        response[:,hh] = Apower[1:n,1:n]*L[:,1]
        response[:,hh] = Apower[1:n,1:n]*response[:,1]
    end

return response

end

function IRFredu_Anatomy_Distr1(A,L,nstep,sh_size)

    np = size(A)[1]
    n = size(A)[2]

    p = Int(np/n)

    lenq = 500 # out of 500 qs, find the one that maximizes the closed form variance
    #Random.seed!(100000*seedindx+10*t+3+seedoffset)
    qgrid = randn(lenq,n)

    for qq = 1:lenq
       qgrid[qq,:] = qgrid[qq,:]./norm(qgrid[qq,:])
    end

    response = zeros(n,nstep)
    Acomp = [A'; sparse(Array(1:n*(p-1)), Array(1:n*(p-1)), ones(n*(p-1)), n*(p-1), np)]

    M = zeros(np, n)
    M[1:n,1:n] = Matrix(1.0I,n,n)

    Var_Distr = zeros(lenq,1)

        for qq = 1:lenq
            temp_q = qgrid[qq,:]
            SIGMAtrq = Symmetric(M*L*temp_q*temp_q'*L'*M')
            SIGMAq = lyapd(Acomp, SIGMAtrq) # solve Lyaponov equation
            Var_alpha = SIGMAq[n_agg+1:end, n_agg+1:end] # variance of alphas
            Var_Distr[qq] = tr(Var_alpha*delta_x*ZZ'*ZZ)
        end

    optimal_q = qgrid[findmax(vec(Var_Distr))[2],:]

    YY_impact = L*optimal_q
    #YY_impact = zeros(n,1) # for ss percentiles
    PhatDensCoef_impact    = coefRecover(YY_impact[n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
    PhatDensNorm_impact    = lnpdfNormalize_unrate(PhatDensCoef_impact,knots,YY_impact[3]+mean_unrate, minimum(xgrid), maximum(xgrid))

    ngrid     = 20
    grid_temp = range(0.2, stop=1.8, length=ngrid);
    grid_temp = [0.0; grid_temp]

    # percentile on impact as sign normalization
    perc_impact = DensPercentiles(PhatDensCoef_impact, knots, YY_impact[3]+mean_unrate, xgrid, grid_temp, 0.1)[1]
    perc_ss = 0.22426370152942293 # 10th percentile
    if (perc_impact - perc_ss) < 0
        optimal_q = -optimal_q
    end

    # Gini on impact as sign normalization
    #PhatDens_impact = pdfEval(xgrid,PhatDensCoef_impact',knots,[PhatDensNorm_impact[1]]);
    #Gini_impact       = GiniCoef(PhatDens_impact, xgrid)
    #Gini_ss = 0.42557515101522825

    #if (Gini_impact - Gini_ss) < 0
    #    optimal_q = -optimal_q
    #end

    response[:,1] = sh_size*L*optimal_q

    Apower = Acomp

    for hh = 2:nstep
        Apower = Apower*Acomp
#        response[:,hh] = Apower[1:n,1:n]*L[:,1]
        response[:,hh] = Apower[1:n,1:n]*response[:,1]
    end

return response

end

function IRFredu_Anatomy_Distr2(A,L,nstep,sh_size)

    np = size(A)[1]
    n = size(A)[2]

    p = Int(np/n)

    lenq = 500 # out of 500 qs, find the one that maximizes the closed form variance
    #Random.seed!(100000*seedindx+10*t+3+seedoffset)
    qgrid = randn(lenq,n)

    for qq = 1:lenq
       qgrid[qq,:] = qgrid[qq,:]./norm(qgrid[qq,:])
    end

    response = zeros(n,nstep)
    Acomp = [A'; sparse(Array(1:n*(p-1)), Array(1:n*(p-1)), ones(n*(p-1)), n*(p-1), np)]

    M = zeros(np, n)
    M[1:n,1:n] = Matrix(1.0I,n,n)

    Var_Distr = zeros(lenq,1)

        for qq = 1:lenq
            temp_q = qgrid[qq,:]
            SIGMAtrq = Symmetric(M*L*temp_q*temp_q'*L'*M')
            SIGMAq = lyapd(Acomp, SIGMAtrq) # solve Lyaponov equation
            Var_alpha = SIGMAq[n_agg+1:end, n_agg+1:end] # variance of alphas
            #Var_Distr[qq] = tr(Var_alpha*delta_x*ZZ'*ZZ)
            ZalphaZ_eig = eigen(ZZ*Var_alpha*ZZ')

            if length(ZalphaZ_eig.values[abs.(ZalphaZ_eig.values).>1/100000000]) > K
                println("More than K non-zero eigenvalues: q_id = $(qq)")
            end

            Var_Distr[qq] = sum(real(ZalphaZ_eig.values)) # complex numbers generated

        end

    optimal_q = qgrid[findmax(vec(Var_Distr))[2],:]

    YY_impact = L*optimal_q
    #YY_impact = zeros(n,1) # for ss percentiles
    PhatDensCoef_impact    = coefRecover(YY_impact[n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
    PhatDensNorm_impact    = lnpdfNormalize_unrate(PhatDensCoef_impact,knots,YY_impact[3]+mean_unrate, minimum(xgrid), maximum(xgrid))

    ngrid     = 20
    grid_temp = range(0.2, stop=1.8, length=ngrid);
    grid_temp = [0.0; grid_temp]

    # percentile on impact as sign normalization
    perc_impact = DensPercentiles(PhatDensCoef_impact, knots, YY_impact[3]+mean_unrate, xgrid, grid_temp, 0.1)[1]
    perc_ss = 0.22426370152942293 # 10th percentile
    if (perc_impact - perc_ss) < 0
        optimal_q = -optimal_q
    end

    # Gini on impact as sign normalization
    #PhatDens_impact = pdfEval(xgrid,PhatDensCoef_impact',knots,[PhatDensNorm_impact[1]]);
    #Gini_impact       = GiniCoef(PhatDens_impact, xgrid)
    #Gini_ss = 0.42557515101522825

    #if (Gini_impact - Gini_ss) < 0
    #    optimal_q = -optimal_q
    #end

    response[:,1] = sh_size*L*optimal_q

    Apower = Acomp

    for hh = 2:nstep
        Apower = Apower*Acomp
#        response[:,hh] = Apower[1:n,1:n]*L[:,1]
        response[:,hh] = Apower[1:n,1:n]*response[:,1]
    end

return response

end




function CompanionForm(B)
    # B is np,n and has equations in columns
    (np,n) = size(B)
    p      = Int(np/n)
    Bprime = B'
    Fmat   = zeros(n*p, n*p)

    for ll = 1:p
        # related to y_{t}
        Fmat[1:n_agg, (ll-1)*n_agg+1:ll*n_agg] = Bprime[1:n_agg,(ll-1)*n+1:(ll-1)*n+n_agg]
        Fmat[1:n_agg, n_agg*p+(ll-1)*n_cross+1:n_agg*p+ll*n_cross] = Bprime[1:n_agg,(ll-1)*n+n_agg+1:ll*n]

        # related to alpha_{t}
        Fmat[n_agg*p+1:n_agg*p+n_cross, (ll-1)*n_agg+1:ll*n_agg] = Bprime[n_agg+1:end,(ll-1)*n+1:(ll-1)*n+n_agg]
        Fmat[n_agg*p+1:n_agg*p+n_cross, n_agg*p+(ll-1)*n_cross+1:n_agg*p+ll*n_cross] = Bprime[n_agg+1:end,(ll-1)*n+n_agg+1:ll*n]
    end

    if p > 1
        # related to other lags of y
        Fmat[n_agg+1:n_agg*p, 1:n_agg*(p-1)] = kron(Matrix(1.0I, p-1, p-1), Matrix(1.0I, n_agg, n_agg))

        # related to other lags of alpha
        Fmat[n_agg*p+n_cross+1:n_agg*p+n_cross*p, n_agg*p+1:n_agg*p+n_cross*(p-1)]  = kron(Matrix(1.0I,p-1,p-1), Matrix(1.0I, n_cross, n_cross))
    end

return Fmat

end
